# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import sys
import os
import time
#from optparse import OptionParser

#TODO: add command line options
#TODO: Add comments + docstrings

#########################################################
class ProcSnapshot(object):
  '''
  Snapshot of the system memory usage.
  '''

  def __init__(self, proc_dir="/proc", chrome_only=False):
    self.proc_dir = proc_dir
    self.chrome_only = chrome_only

    self.timestamp = time.time()

    proc_list = self.getProcList()
    self.chrome_pids_list = sorted(proc_list[0], key=int)
    self.sys_pids_list = sorted(proc_list[1], key=int)
    self.getSmaps()

    self.chrome_types = {}
    self.meminfo = MemInfo(self.proc_dir)
    self.gpu_mem = GPUGraphicMem()
    self.getChromeTypes()

  def getProcList(self):
    '''
    Return tuple with two lists:
      - PIDs of all Chrome Procs
      - PIDs of non Chrome Procs
    Arguments:
      - proc_dir: Path to directory containing the processes information
    '''
    chrome_pids = []
    sys_pids = []
    # List of all dirs and files in /proc
    for item in os.listdir(self.proc_dir):
      filepath = os.path.join(self.proc_dir, item)
      #if item == PID
      if os.path.isdir(filepath) and item.isdigit():
        try:
          f = open(filepath + "/comm", 'r')
          if self._isChrome(f):
            chrome_pids.append(item)
          else:
            sys_pids.append(item)
        except IOError: pass
        f.close()

    return chrome_pids, sys_pids

  def _isChrome(self, f):
    '''Return True if file is a chrome process. False otherwise'''
    s = f.readline().strip()
    return (s.startswith("chrome"))

  def getChromeTypes(self):
    foundBrowser = False
    for pid in self.chrome_pids_list:
      pid_type = self.getSingleChromeType(pid)
      if pid_type != "chrome":
        self.chrome_types[pid] = pid_type
      elif not foundBrowser:
      # self.chrome_pids_list is sorted and browser PID < sandbox PID.
      # This is a hack, as there is no way of figuring out what is the chrome
      # browser and what is the chrome sandbox helper outside chrome.
        self.chrome_types[pid] = "browser"
        foundBrowser = True
      else:
        self.chrome_types[pid] = "sandbox"

  def getSingleChromeType(self, pid):
    if pid == "Total":
      return ""
    try:
      filename = os.path.join(self.proc_dir, pid, "cmdline")
      f = open(filename, 'r')
      cmdline = f.readline()
      f.close()
    except IOError:
      return "Process died during run time"

    index = cmdline.find("--type=")
    if index != -1:
      # sample: /opt/google/chrome/chrome[...] --type:[...] --[...]
      # special case sample:
      #/opt/google/chrome/chrome-sandbox/opt/google/chrome/chrome--type=zygote\
      #--log-level=1 ==> this requires another split at "--", as it has no
      # spaces
      pid_type = cmdline[index:].split()[0].split("--")[1]
      pid_type = pid_type[5:]
    else:
      # process is either the browser or sandbox
      pid_type = "chrome"

    return pid_type

  def getSmaps(self):
    self.chrome_smaps = Smaps(self.chrome_pids_list, self.proc_dir)
    if not self.chrome_only:
      self.sys_smaps = Smaps(self.sys_pids_list, self.proc_dir)

  def totalOutput(self):
    totalOutput(chrome_pid_dict, "Chrome")
    if not self.chrome_only:
      totalOutput(sys_pid_dict, "Non-Chrome")
    #for mem_type in self.
    #print "Mem Type: {0:7}".format()

  def sortedOutput(self, mem_type, reverse=False):
    mem_type = mem_type.capitalize()
    self._headerOutput("Chrome", mem_type, self.chrome_types)
    chrome_sorted_list = self.chrome_smaps.sortByMemType(mem_type, reverse)
    self.sortedOutputHelper(chrome_sorted_list, mem_type, self.chrome_types)

    if not self.chrome_only:
      print ""
      self._headerOutput("Non-Chrome", mem_type)
      sys_sorted_list = self.sys_smaps.sortByMemType(mem_type, reverse)
      self.sortedOutputHelper(sys_sorted_list, mem_type)

  def _headerOutput(self, proc_type, mem_type, types_dict={}):
    print proc_type, "processes", mem_type, "memory  usage, in kBs:\n"
    print "PID".rjust(5), mem_type.rjust(12),
    if self.meminfo != None:
      print "Total %".rjust(12),
    if types_dict != {}:
      print "Type".rjust(12)
    else:
      print ""

  def sortedOutputHelper(self, sorted_list, mem_type, proc_types={}):

    for item in sorted_list:
      mem_size = item[1][mem_type]
      pid = item[0]
      print pid.rjust(5),
      print repr(mem_size).rjust(12),
      if self.meminfo != None:
        percent_total_size = round(100.0 * mem_size /
                              self.meminfo.memtotal, 5)
        print "{0:f}".format(percent_total_size)[:5].rjust(12),
      if proc_types != {} and pid != "Total":
        print proc_types[pid].rjust(12)
      else:
        print ""

  def sytemOutput(self):
    pass

  def GPUOutput(self):
    if self.gpu_mem != None:
      print "\nGPU memory usage, in kBs"
      print "Current GPU memory use:", self.gpu_mem.current
      print "Total shared memory belonging to GPU:", self.gpu_mem.total

  def createGraph(self, dest):
    pass

########################################################
class Smaps(object):

  def __init__(self, pids_list, proc_dir="/proc"):
    self.exec_mem = {}
    self.exec_mem["Total"] = {}
    smaps = self.getSmapsDict(pids_list, proc_dir)
    self.pid = smaps[0]
    #self.page_addr = smaps[1]
    self.valid_mem_types = self.getValidMemTypes()

  def getSmapsDict(self, pids_list, proc_dir):
    '''
    Return tuple with two filled dictionary containing information taken from
    processes smaps file inside proc_dir. Dictionaries are, respectively:
      - {PID:{memory type:size}}
      - {page address:{PID:{memory type:size}}}

    Arguments:
      - pids_list: List of strings containing the PIDs.
      - proc_dir: Path to directory containing the processes information
    '''
    page_addr_dict = {}
    pid_dict = {}
    pid_dict["Total"] = {}

    for pid in pids_list[:]:
      filename = os.path.join(proc_dir, pid, "smaps")
      mem_type_dict = {}

      try:
        f = open(filename, 'r')
        f_lines = f.readlines()
        f.close()
        for line in f_lines:
          line = line.split()

          # Check if line is referent to page address
          if line[2] != "kB":
            page_addr = line[0]
            if not page_addr_dict.has_key(page_addr):
              page_addr_dict[page_addr] = {}

            if 'x' in line[1]:
              is_exec = True
            else:
              is_exec = False

          # If line is referent to memory
          else:
            mem_type = line[0].strip(':')
            mem_size = int(line[1])
            if mem_type not in ("KernelPageSize", "MMUPageSize"):
              self.addToPageAddrDict(page_addr_dict, page_addr, pid, mem_type,
                                mem_size)
              self.addToMemTypeDict(mem_type_dict, mem_type, mem_size)
              if is_exec:
                #self.addToExecDict(pid, mem_type, mem_size)
                pass

        # Add PID memory to Total
        for item in mem_type_dict:
          pid_dict["Total"][item] = (mem_type_dict[item] +
                                   pid_dict["Total"].get(item, 0))

        # Account for kernel processes (empty smaps)
        if mem_type_dict !=  {}:
          pid_dict[pid] = mem_type_dict
        else:
          pids_list.remove(pid)
          pass

      except IOError:
        pids_list.remove(pid)

    return pid_dict, page_addr_dict

  def addToPageAddrDict(self, page_addr_dict, page_addr, pid, mem_type,
                        mem_size):
    if not page_addr_dict[page_addr].has_key(pid):
      page_addr_dict[page_addr][pid] = {}
    page_addr_dict[page_addr][pid][mem_type] = mem_size


  def addToMemTypeDict(self, mem_type_dict, mem_type, mem_size):
    mem_type_dict[mem_type] = mem_type_dict.get(mem_type, 0) + mem_size
    for item in ("Private", "Shared"):
      if mem_type.startswith(item):
        mem_type_dict[item] = mem_type_dict.get(item, 0) + mem_size

  def addToExecDict(self, pid, mem_type, mem_size):
    '''
    Add executable memory size to self.exec_mem
    '''
    if pid not in self.exec_mem:
      self.exec_mem[pid] = {}

    self.exec_mem[pid][mem_type] = (mem_size +
                                  self.exec_mem[pid].get(mem_type, 0))
    self.exec_mem["Total"][mem_type] = (mem_size +
                                  self.exec_mem["Total"].get(mem_type, 0))


  def sortByMemType(self, mem_type, reverse_=False):
    self.pid_list = self.pid.items()
    self.pid_list.sort(key=lambda pid: pid[1][mem_type],
                       reverse=reverse_)
    return self.pid_list

  def getValidMemTypes(self):
    valid_mem_types = set()
    for mem_type in self.pid["Total"].keys():
      valid_mem_types.add(mem_type.upper())
      valid_mem_types.add(mem_type.lower())
      valid_mem_types.add(mem_type)
    return valid_mem_types

###############################################################
class MemInfo(object):

  def __init__(self, proc_dir="/proc"):
    self.meminfo = self.getMemInfo(proc_dir)
    self.memtotal = self.meminfo["MemTotal"]
    #TODO(thiagog):
    #self.total
    #self.available
    # other interesting values that we would like to consider

  def getMemInfo(self, source):
    '''
    Return a dictionary with all the infomation contained in "source"/meminfo.
    '''
    filename = os.path.join(source, "meminfo")
    f = open(filename)
    f_lines = f.readlines()
    f.close()

    meminfo_dict = {}
    for line in f_lines:
      line = line.split()
      meminfo_dict[line[0].strip(":")] = int(line[1])

    return meminfo_dict

#########################################################
class GPUGraphicMem(object):
  def __init__(self):
    self.mem = self.getGPUMem()
    if self.mem != None:
      self.total = self.mem[0]
      self.current = self.mem[1]

  def getGPUMem(self):
    '''
    Return a tuple with total memory being allocated for GPU use and
    total memory it is currently using, in kBs
    NOTE: only for ChromeOS
    '''
    filename = "/sys/kernel/debug/dri/0/i915_gem_objects"
    if os.path.exists(filename):
      f = open(filename, 'r')
      f_lines = f.readlines()
      f.close()
      # Info on first line, 3rd element, in bytes
      GPU_total = int(f_lines[0].split()[2]) / 1024
      # Info on second line, 5th element, inside brackets, in bytes
      GPU_current = int(f_lines[1].split()[4][1:-1]) / 1024
      return GPU_total, GPU_current
    else:
      return None


##########################################################
class SelfMemUse(object):
  '''
  Objects stores how much memory the current process is using. Primarily used
  for monitoring the script's memory consumption, and to verify how the script
  is affect the system memory.
  '''
  def __init__(self):
   self.pid = os.getpid()

  def getMem(self):
    smaps = {}
    smaps_path = "/proc/self/smaps"
    f = open(smaps_path, 'r')
    f_lines = f.readlines()
    f.close()
    for line in f_lines:
      line = line.split()

     #Check if line is referent to page address
      if line[2] == "kB":
        mem_type = line[0].strip(':')
        mem_size = int(line[1])
        if mem_type not in ("KernelPageSize", "MMUPageSize"):
          smaps[mem_type] = smaps.get(mem_type, 0) + mem_size
          for item in ("Private", "Shared"):
            if mem_type.startswith(item):
              smaps[item] = smaps.get(mem_type, 0) + mem_size

    return smaps

  mem = property(getMem, doc="Memory consumption of script, from smaps")


##########################################################
if __name__ == "__main__":
  snapshot = ProcSnapshot()
  self_mem = SelfMemUse()
  #print snapshot.chrome_pids_list, '\n'
  #print snapshot.sys_pids_list, '\n'

  #Testing if code handles wrong sorting input
  if "lalala" in snapshot.chrome_smaps.valid_mem_types:
    snapshot.sortedOutput("lalala")

  elif "PSS" in snapshot.chrome_smaps.valid_mem_types:
    snapshot.sortedOutput("PSS")
    print self_mem.pid
    print self_mem.mem
